import matplotlib.pyplot as plt
import numpy as np


# 假设这里复用之前定义的DASNN类中的一些参数逻辑等，简化处理，只提取关键结构绘制示意
class DASNN:
    def __init__(self, shape_mode="squaresmall"):
        self.shape_mode = shape_mode
        if self.shape_mode == "squaresmall":
            self.nb_steps = 5
            self.nb_channels = 32
            self.kernels = 5
            self.strides = 3
            self.outconvsize = 9
        elif self.shape_mode == "squarebig":
            self.nb_steps = 7
            self.nb_channels = 16
            self.kernels = 4
            self.strides = 2
            self.outconvsize = 49
        elif self.shape_mode == "rectsame":
            self.nb_steps = 4
            self.nb_channels = 64
            self.kernels = 4
            self.strides = 2
            self.outconvsize = 56
        elif self.shape_mode == "rectsamemaxpool":
            self.nb_steps = 4
            self.nb_channels = 64
            self.kernels = 4
            self.strides = 2
            self.outconvsize = 7
            self.maxpool = True
        elif self.shape_mode == "rectplus":
            self.nb_steps = 5
            self.nb_channels = 64
            self.kernels = (10, 3)
            self.strides = (4, 2)
            self.outconvsize = 1
        else:
            self.nb_steps = 4
            self.nb_channels = 64
            self.kernels = (10, 3)
            self.strides = (4, 2)
            self.outconvsize = 8


# 实例化模型（这里随便选个shape_mode示例，用于获取对应的层数等信息来绘制，可按需改）
model = DASNN(shape_mode="squaresmall")

# 设置图形大小
plt.figure(figsize=(8, 12))

# 计算各层的垂直位置（均匀分布示例）
layer_y_positions = np.linspace(0.95, 0.05, model.nb_steps + 2)  # +2是算上输入层和输出层

# 绘制输入层矩形，添加文字标注及输入尺寸示意（这里假设输入为单通道图像，尺寸示例为900x900，可按需改）
input_layer = plt.Rectangle((0.1, 0.95), 0.2, 0.1, facecolor='lightblue', label='Input Layer\n(1x900x900 Image)')
plt.gca().add_patch(input_layer)
plt.text(0.2, 0.99, "Input", ha='center', va='center', fontsize=10)

# 绘制卷积层矩形（纵向排列多个卷积层，展示更多细节）
for i in range(model.nb_steps):
    conv_layer = plt.Rectangle((0.3, layer_y_positions[i + 1]), 0.3, 0.08, facecolor='orange',
                               label='Conv Layer' if i == 0 else "")
    plt.gca().add_patch(conv_layer)
    plt.text(0.45, layer_y_positions[i + 1] + 0.04, f"Conv {i + 1}", ha='center', va='center', fontsize=10)
    # 展示卷积层细节参数（示例展示部分关键参数，可进一步完善）
    detail_text = f"Kernel: {model.kernels}\nStride: {model.strides}\nChannels: {model.nb_channels * 2 ** (i + 1)}"
    if hasattr(model,'maxpool') and model.maxpool and i < model.nb_steps - 1:
        detail_text += "\nMaxPool: Yes"
    plt.text(0.45, layer_y_positions[i + 1] - 0.02, detail_text, ha='center', va='center', fontsize=8,
             color='white')
    if i < model.nb_steps - 1:
        plt.arrow(0.6, layer_y_positions[i + 1] + 0.04, 0, -0.08, head_width=0.02, head_length=0.05, fc='black',
                  ec='black')

# 绘制全连接层矩形（如果有），展示全连接层细节（示例展示节点数等简单信息，可完善）
if model.nb_steps > 0:
    if model.outFC > 0:
        intermediate = (model.outconvsize * model.nb_channels * 2 ** (model.nb_steps) + model.outsize) // 2
        fc_layer = plt.Rectangle((0.8, 0.3), 0.4, 0.1, facecolor='green', label='FC Layer(s)')
        plt.gca().add_patch(fc_layer)
        plt.text(1.0, 0.35, "FC", ha='center', va='center', fontsize=10)
        detail_text = f"Input Nodes: {model.outconvsize * model.nb_channels * 2 ** (model.nb_steps)}\n" \
                      f"Intermediate Nodes: {intermediate}\nOutput Nodes: {model.outsize}"
        plt.text(1.0, 0.3 - 0.02, detail_text, ha='center', va='center', fontsize=8, color='white')
    else:
        fc_layer = plt.Rectangle((0.8, 0.3), 0.4, 0.1, facecolor='green', label='FC Layer')
        plt.gca().add_patch(fc_layer)
        plt.text(1.0, 0.35, "FC", ha='center', va='center', fontsize=10)
        detail_text = f"Input Nodes: {model.outconvsize * model.nb_channels * 2 ** (model.nb_steps)}\n" \
                      f"Output Nodes: {model.outsize}"
        plt.text(1.0, 0.3 - 0.02, detail_text, ha='center', va='center', fontsize=8, color='white')
    plt.arrow(0.6, layer_y_positions[-1] + 0.04, 0, -0.16, head_width=0.02, head_length=0.05, fc='black', ec='black')

# 绘制输出层矩形，添加文字标注及输出尺寸示意（示例输出维度为1x3，可按需改）
output_layer = plt.Rectangle((1.2, 0.1), 0.2, 0.1, facecolor='lightgreen', label='Output Layer\n(1x3 Prediction)')
plt.gca().add_patch(output_layer)
plt.text(1.3, 0.15, "Output", ha='center', va='center', fontsize=10)


#进行测试修改

# 设置坐标轴范围等
plt.xlim(0, 1.5)
plt.ylim(0, 1)
plt.axis('off')
plt.title("DASNN Network Model Structure")
plt.legend(loc='upper center', bbox_to_anchor=(0.5, 1.15), ncol=4, fancybox=True, shadow=True)
plt.show()